from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension


setup(
    name="only_costh", 
    version="1.2.3", 
    description="norm_clamp_div optimization, without matrices multiplication", 
    long_description="2nd topic of Spark project, optimization of `norm + clamp + div`, without the PyTorch C++ API `torch::mm`", 
    url="https://lixiaobao.niubility.cn/", 
    author="Xiaobao Li", 
    author_email="miznchimaki@gmail.com", 
    ext_modules=[
        CUDAExtension("only_costh", [
            "costh_cuda.cpp", 
            "costh_cuda_kernel.cu"
        ])
    ], 
    cmdclass={
        "build_ext": BuildExtension
    })
